from typing import Any, Callable, Dict, List, Optional, Set, Union
from collections import deque

from ezrobot.utils import Timer

LabelType = Union[str, int]


class GraphNode:

    __slots__ = {'name', 'func', 'dependency', 'result', 'init_val'}
    
    def __init__(self, name: LabelType, func: Callable, dependency: List['GraphNode'] = None, init: Any = None):
        self.name = name
        self.func = func
        self.dependency = dependency or []
        self.init_val = init
        self.result = init

    def reset(self) -> None:
        """重置节点状态，清除执行结果"""
        self.result = self.init_val
    
    def __repr__(self):
        return f"GraphNode(name={self.name}, dependency={[dep.name for dep in self.dependency]})"


class GraphRunner:
    """管理和执行计算图"""
    
    def __init__(self):
        """初始化GraphRunner"""
        self._nodes: Dict[LabelType, GraphNode] = {}  # 存储所有节点的字典，键为节点名称，值为Node对象
        self._childrens: Dict[LabelType, Set[LabelType]] = {}  # 存储每个节点的子节点集合
        self._run_order: List[LabelType] = None  # 缓存拓扑排序结果
        self._timer = None  # 用于异步运行的定时器
        self._last_node_name = None
        
    def add_node(self, name: LabelType, func: Callable, dependency=None, init=None) -> 'GraphRunner':
        """
        添加一个节点到计算图中, 该节点可以有多个依赖节点，依赖节点的结果将作为当前节点函数的输入参数。

        Parameters:
        - name : LabelType, 唯一节点名称, 可为 >= 0 的整数或字符串
        - func : Callable
        - dependency : List[LabelType] | LabelType, default None 表示无依赖, -1 表示依赖上一个节点
        - init : Any, (可选) 初始化值，用于缓存节点的初始状态, 表示该节点是一个环节点
        """
        # 检查是否有重复的节点名称
        if name in self._nodes:
            raise ValueError(f"Node name '{name}' already exists")
            
        # 查找依赖节点
        dep_nodes = []
        if dependency == -1:
            dep_names = [self._last_node_name] if self._last_node_name is not None else []
        elif isinstance(dependency, (str, int)):
            dep_names = [dependency]
        else:
            dep_names = dependency if isinstance(dependency, list) else []
            
        for dep_name in dep_names:
            assert dep_name in self._nodes, f"Dependency '{dep_name}' does not exist in the graph"
            dep_nodes.append(self._nodes[dep_name])
                
        # 创建并添加节点
        new_node = GraphNode(name, func, dep_nodes, init)
        if init is not None:  # 环节点
            new_node.dependency = [new_node]
        self._nodes[name] = new_node
        
        # 更新依赖节点的子节点列表
        self._childrens[name] = set()  # 初始化子节点集合
        for dep_name in dep_names:
            if dep_name not in self._childrens:
                self._childrens[dep_name] = set(name)
            else:
                self._childrens[dep_name].add(name)

        self._run_order = None  # 清除之前的运行顺序缓存
        self._last_node_name = name  # 更新最后一个节点名称
        return self
        
    def get_node(self, name: LabelType) -> Optional[GraphNode]:
        return self._nodes.get(name, None)
    
    def reset(self) -> None:
        """重置所有节点状态"""
        for node in self._nodes.values():
            node.reset()
    
    def _get_topological_order(self) -> List[str]:
        """
        返回图的拓扑排序结果
        """
        # 计算每个节点的入度
        in_degree = {}
        for name, node in self._nodes.items():
            if len(node.dependency) == 1 and node.dependency[0].name == name:  # 环节点
                in_degree[name] = 0
            else:
                in_degree[name] = len(node.dependency)

        # 初始化队列，将入度为0的节点加入队列
        queue = deque([name for name, degree in in_degree.items() if degree == 0])
        result = []
        
        while queue:
            name = queue.popleft()
            result.append(name)
            # 遍历子节点，减少它们的入度
            for child_name in self._childrens[name]:
                in_degree[child_name] -= 1
                if in_degree[child_name] == 0:
                    queue.append(child_name)

        if len(result) != len(self._nodes):
            raise ValueError("Graph contains a cycle.")
            
        return result

    def run_once(self):
        if self._run_order is None:
            self._run_order = self._get_topological_order()
        
        for node_name in self._run_order:
            node = self._nodes[node_name]
            # 收集依赖节点的结果
            dependencies_results = [dep.result for dep in node.dependency]
            # 执行当前节点的函数
            try:
                node.result = node.func(*dependencies_results)
            except Exception as e:
                print(f"Error running node '{node_name}': {e}")
                raise e
            
        return self._nodes[self._run_order[-1]].result  # 返回最后一个节点的结果
        
    def run_async(self, frequency) -> 'GraphRunner':
        """
        异步运行计算图
        """
        if self._timer is None:
            self._timer = Timer(frequency, self.run_once, start=False)
            self._timer.start()
        return self
    
    def stop(self):
        """
        停止异步运行
        """
        if self._timer is not None:
            self._timer.stop()
            self._timer = None
    
    def run(self, frequency):
        """
        同步运行计算图
        """
        if self._timer is None:
            self._timer = Timer(frequency, self.run_once, start=False)
            self._timer.run()

    def __repr__(self):
        ret = "GraphRunner:\n"
        for name, node in self._nodes.items():
            ret += f"- {str(node)}\n"
        return ret[:-1]



# 使用示例
if __name__ == "__main__":
    import time
    
    # 创建一个GraphRunner实例
    runner = (
        GraphRunner()
        .add_node("a", lambda x: x+1, None, init=0)
        .add_node("scale", lambda: 2, None)
        .add_node("c", lambda a, s: a * s, dependency=["a", "scale"])
        .add_node("print", print, -1)
    )
    
    print(runner._get_topological_order())
    
    runner.run_async(frequency=1)

    time.sleep(5)
    